# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: xguo

import datetime
import json

import matplotlib.pyplot as plt
import pandas as pd

from absl import app, flags

flags.DEFINE_string('config', None, 'config file.')
flags.DEFINE_string('oe', None, 'order event file.')


def get_hs_time(timestamp):
  return datetime.datetime.fromtimestamp(timestamp / 1e9).strftime('%Y%m%d %H:%M:%S.%f')


def save_to_file(fn, obj):
  with open(fn, 'w') as ofile:
    json.dump(obj, ofile, indent=2)


def extract_balance(json_log, currency):
  res = []
  for each in json_log:
    if each['data']['type'] == 'BALANCE':
      ts = each['timestamp']
      balance_info = each['data']['balance']
      if 'each_balance' in balance_info:
        for pos in balance_info['each_balance']:
          if pos['currency'] == currency:
            data = pos
            break
      else:
        print(each)
        continue

      info = {
          'symbol': data['currency'],
          'timestamp': ts,
          'hs_timestamp': get_hs_time(ts),
          'net_position': data['total'],
      }
      res.append(info)
  return sorted(res, key=lambda x: x['timestamp'])


def extract_position(json_log, product):
  res = []
  for each in json_log:
    if each['data']['type'] == 'POSITION':
      ts = each['timestamp']
      position_info = each['data']['position']
      data = None
      if 'each_position' in position_info:
        for pos in position_info['each_position']:
          if pos['symbol'] == product:
            data = pos
            break
      else:
        print(each)
        continue

      if data is not None:
        info = {
            'symbol': data['symbol'],
            'timestamp': ts,
            'hs_timestamp': get_hs_time(ts),
            'net_position': data['net_position'],
        }
        res.append(info)
  return sorted(res, key=lambda x: x['timestamp'])


def get_sign(side):
  if 'BUY' in side:
    return +1
  elif 'SELL' in side:
    return -1
  else:
    raise ValueError()


def extract_fill(json_log, product):
  known_fill_id = set()
  res = []
  for each in json_log:
    ts = each['timestamp']
    if each['data']['type'] != 'ORDER_EVENT':
      continue

    if 'oe' in each['data']:
      oe = each['data']['oe']
    else:
      oe = each['data']['event']

    if oe['type'] != "ORDER_FILLED":
      continue

    if oe['symbol'] != product:
      continue

    if oe['fill_id'] in known_fill_id:
      print('dup fill')
      continue
    else:
      known_fill_id.add(oe['fill_id'])

    info = {
        'symbol': oe['symbol'],
        'price': oe['fill_price'],
        'side': oe['order_side'],
        'qty': oe['fill_qty'],
        'sqty': oe['fill_qty'] * get_sign(oe['order_side']),
        'timestamp': ts,
        'hs_timestamp': get_hs_time(ts),
    }
    res.append(info)
  return sorted(res, key=lambda x: x['timestamp'])


def main(argv):
  with open(flags.FLAGS.config) as infile:
    config = json.load(infile)

  with open(flags.FLAGS.oe) as infile:
    data = json.load(infile)

  product = config['symbol']
  if 'Spot' in config['mea']:
    currency = config['symbol'].split('-')[0]
    col = extract_balance(data, currency)
  else:
    col = extract_position(data, product)

  # col = [x for x in col if x['timestamp'] >= fix_time]
  save_to_file('pos.json', col)
  print(len(col))
  df = pd.DataFrame(col[4:])
  df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
  plt.plot(df.timestamp, df.net_position, marker='.', label='Position')
  init_qty = df.net_position[0]
  print(init_qty)
  # init_qty -= 1

  col = extract_fill(data, product)
  # col = [x for x in col if x['timestamp'] >= fix_time]
  save_to_file('fill.json', col)

  df = pd.DataFrame(col)
  try:
    df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
  except Exception:
    import pdb
    pdb.set_trace()

  r = df.sqty.cumsum() + init_qty
  plt.plot(df.timestamp, r, marker='.', label='Position derived from fill message')
  print("qty fill_pq:", df.qty.sum(), (df.qty * df.price).sum())

  plt.title(config['symbol'])
  plt.grid(True)
  plt.legend()
  plt.savefig("position.png")
  plt.show()


if __name__ == "__main__":
  app.run(main)
